#include <cmath>
#include <iostream>
#include <sstream>
#include <string>
#include "hHuffman.h"

void HuffmanCode::genTree()
{
	std::multimap<size_t, _hHuffmanDetail::Node*> freqTab;
	for (auto& [ch, freq] : _charTab)
	{
		if (!freq)
			continue;

		auto pNode = new _hHuffmanDetail::CharNode;
		pNode->_ch = ch;
		pNode->_freq = freq;
		freqTab.insert(std::make_pair(freq, pNode));
	}

	if (freqTab.empty())
		return;

	size_t n = freqTab.size() - 1;
	for (size_t i = 0; i < n; ++i)
	{
		auto pNode = new _hHuffmanDetail::Node;
		pNode->_left = freqTab.begin()->second;
		freqTab.erase(freqTab.begin());
		pNode->_right = freqTab.begin()->second;
		freqTab.erase(freqTab.begin());
		pNode->_freq = pNode->_left->_freq + pNode->_right->_freq;
		freqTab.insert(std::make_pair(pNode->_freq, pNode));
	}

	_pTreeNode = freqTab.begin()->second;
}

struct FillOsTreeParam
{
	enum NeedBitType
	{
		NeedBitType_Left = 0x01,
		NeedBitType_Right = 0x02,
	};
	size_t _nilNum = 3;
	size_t formatNum1 = 2;
	size_t formatNum2 = 3;
	size_t formatNum = 1;
	size_t _heightMax = 0;
	size_t _nodeMax = 0;
	size_t _curPos = 0;
	std::map<size_t, std::map<size_t, std::pair<size_t, size_t>>> _dataPos;
	std::map<size_t, std::map<size_t, std::string>> _dataStr;

	void fillHeightMaxData(_hHuffmanDetail::Node* pNode);
	void fillAllByHeightMaxData(const _hHuffmanDetail::Node* pNode);
	std::ostream& fillOs(std::ostream& os) const;
	size_t fillTreeData(const _hHuffmanDetail::Node* src, std::string& dst, size_t len = 0, size_t needFlag = 0);
	void freshFormatNum();
};

void HuffmanCode::showTree()
{
	FillOsTreeParam param;
	param._heightMax = _pTreeNode->calcHeight();
	param._nodeMax = (size_t)exp2(param._heightMax);
	param.fillHeightMaxData(_pTreeNode);
	param.fillAllByHeightMaxData(_pTreeNode);
	param.fillOs(std::cout);
}

void FillOsTreeParam::fillHeightMaxData(_hHuffmanDetail::Node* pNode)
{
	if (!pNode)
		return;

	for (size_t i = 0; i < _nodeMax; ++i)
	{
		auto& posPair = _dataPos[_heightMax][i];
		posPair.first = _curPos;

		const _hHuffmanDetail::Node* pChildNode = pNode->getNodeByHeight(_heightMax, i);
		if (pChildNode)
		{
			posPair.second = fillTreeData(pChildNode, _dataStr[_heightMax][_curPos]);
			posPair.first = (2 * _curPos + posPair.second) / 2;
			_curPos += posPair.second + formatNum;
		}
		else
		{
			posPair.first = (2 * _curPos + _nilNum) / 2;
			_curPos += _nilNum + formatNum;
		}

		freshFormatNum();
	}
}

void FillOsTreeParam::fillAllByHeightMaxData(const _hHuffmanDetail::Node* pNode)
{
	if (!pNode)
		return;

	while (true)
	{
		if (!_heightMax)
		{
			auto& rootPosMap = _dataPos[0];
			if (!rootPosMap.empty())
				return;

			const _hHuffmanDetail::Node* pChildNode = pNode->getNodeByHeight(0, 0);
			if (!pChildNode)
				return;

			auto& rootPosPair = rootPosMap[0];
			auto& leftPosPair = _dataPos[1][0];
			auto& rightPosPair = _dataPos[1][1];
			rootPosPair.first = (leftPosPair.first + rightPosPair.first) / 2;
			rootPosPair.second = fillTreeData(
				pChildNode,
				_dataStr[0][leftPosPair.first],
				rightPosPair.first + 1 - leftPosPair.first,
				(rightPosPair.second ? NeedBitType_Right : 0) | (leftPosPair.second ? NeedBitType_Left : 0));
			return;
		}

		auto& parentPos = _dataPos[_heightMax - 1];
		auto& parentStr = _dataStr[_heightMax - 1];
		_nodeMax /= 2;
		for (size_t i = 0; i < _nodeMax; ++i)
		{
			auto& leftPosPair = _dataPos[_heightMax][2 * i];
			auto& rightPosPair = _dataPos[_heightMax][2 * i + 1];
			parentPos[i].first = (leftPosPair.first + rightPosPair.first) / 2;
			const _hHuffmanDetail::Node* pChildNode = pNode->getNodeByHeight(_heightMax - 1, i);
			if (pChildNode)
				parentPos[i].second = fillTreeData(
					pChildNode,
					parentStr[leftPosPair.first],
					rightPosPair.first + 1 - leftPosPair.first,
					(rightPosPair.second ? NeedBitType_Right : 0) | (leftPosPair.second ? NeedBitType_Left : 0));
		}

		--_heightMax;
	}
}

size_t FillOsTreeParam::fillTreeData(const _hHuffmanDetail::Node* src, std::string& dst, size_t len, size_t needFlag)
{
	std::ostringstream tmpDataOs;
	if (src)
		src->fillOs(tmpDataOs);
	else
		tmpDataOs << "";

	size_t dataLen = tmpDataOs.str().size();
	if (len <= 2 + dataLen)
	{
		dst = tmpDataOs.str();
		return dataLen;
	}

	size_t halfLen = (len - dataLen) / 2;
	std::ostringstream tmpOs;

	if (needFlag & NeedBitType_Left)
		tmpOs << "+" << std::string(halfLen - 1, '-');
	else
		tmpOs << std::string(halfLen, ' ');

	tmpOs << tmpDataOs.str();

	if (needFlag & NeedBitType_Right)
		tmpOs << std::string(len - dataLen - halfLen - 1, '-') << "+";
	else
		tmpOs << std::string(len - dataLen - halfLen, ' ');

	dst = tmpOs.str();
	return dst.size();
}

std::ostream& FillOsTreeParam::fillOs(std::ostream& os) const
{
	for (auto& nodeMap : _dataStr)
	{
		bool hasChild = false;
		std::ostringstream tmpOs;
		size_t pos = 0;
		for (auto& node : nodeMap.second)
		{
			os << std::string(node.first - pos, ' ');
			tmpOs << std::string(node.first - pos, ' ');
			os << node.second;
			for (auto ch : node.second)
			{
				if ('+' == ch)
				{
					tmpOs << '|';
					hasChild = true;
				}
				else
					tmpOs << ' ';
			}

			pos = node.first + node.second.size();
		}

		os << std::endl;
		if (hasChild)
			os << tmpOs.str() << std::endl;
	}

	return os;
}

void FillOsTreeParam::freshFormatNum()
{
	if (formatNum == formatNum1)
		formatNum = formatNum2;
	else
		formatNum = formatNum1;
}